#include <iostream>
#include <vector>
#include <string>
#include <map>
#include <functional>
#include <cstdint>
#include <sstream>
#include <algorithm>
#include <stdexcept>

using namespace std;

// 常量定义
const int MEM_LEN = 512;
const int CODE_SEG_BEGIN = 0;
const int CODE_SEG_LEN = 256;
const int DATA_SEG_BEGIN = 256;
const int DATA_SEG_LEN = 256;
const int COINS_REMAIN_INIT = 2618;

// 全局变量
vector<int64_t> MEM(MEM_LEN, 0);
int PC = 0;
vector<int64_t> REGSTACK;
int CYCLES = 0;
int COINS_REMAIN = COINS_REMAIN_INIT;

// 指令定义
struct Instruction {
    string name;
    int code;
    int args;
    function<void()> action;
};

vector<Instruction> ISA = {
    {"NOP", 0, 0, []() { PC += 1; }},
    {"PUSH", 1, 1, []() { 
        int arg = MEM[CODE_SEG_BEGIN + PC + 1];
        REGSTACK.push_back(arg);
        PC += 2;
    }},
    {"POP", 2, 0, []() { 
        if (!REGSTACK.empty()) {
            REGSTACK.pop_back(); 
        }
        PC += 1;
    }},
    {"ADD", 11, 0, []() { 
        if (REGSTACK.size() < 2) return;
        int64_t op2 = REGSTACK.back(); REGSTACK.pop_back();
        int64_t op1 = REGSTACK.back(); REGSTACK.pop_back();
        REGSTACK.push_back(op1 + op2);
        PC += 1;
    }},
    {"SUB", 12, 0, []() { 
        if (REGSTACK.size() < 2) return;
        int64_t op2 = REGSTACK.back(); REGSTACK.pop_back();
        int64_t op1 = REGSTACK.back(); REGSTACK.pop_back();
        REGSTACK.push_back(op1 - op2);
        PC += 1;
    }},
    {"MUL", 13, 0, []() { 
        if (REGSTACK.size() < 2) return;
        int64_t op2 = REGSTACK.back(); REGSTACK.pop_back();
        int64_t op1 = REGSTACK.back(); REGSTACK.pop_back();
        REGSTACK.push_back(op1 * op2);
        PC += 1;
    }},
    {"DIV", 14, 0, []() { 
        if (REGSTACK.size() < 2) return;
        int64_t op2 = REGSTACK.back(); REGSTACK.pop_back();
        int64_t op1 = REGSTACK.back(); REGSTACK.pop_back();
        if (op2 != 0) {
            REGSTACK.push_back(op1 / op2);
        } else {
            REGSTACK.push_back(0);
        }
        PC += 1;
    }},
    {"MOD", 15, 0, []() { 
        if (REGSTACK.size() < 2) return;
        int64_t op2 = REGSTACK.back(); REGSTACK.pop_back();
        int64_t op1 = REGSTACK.back(); REGSTACK.pop_back();
        if (op2 != 0) {
            REGSTACK.push_back(op1 % op2);
        } else {
            REGSTACK.push_back(0);
        }
        PC += 1;
    }},
    {"AND", 21, 0, []() { 
        if (REGSTACK.size() < 2) return;
        int64_t op2 = REGSTACK.back(); REGSTACK.pop_back();
        int64_t op1 = REGSTACK.back(); REGSTACK.pop_back();
        REGSTACK.push_back(op1 & op2);
        PC += 1;
    }},
    {"OR", 22, 0, []() { 
        if (REGSTACK.size() < 2) return;
        int64_t op2 = REGSTACK.back(); REGSTACK.pop_back();
        int64_t op1 = REGSTACK.back(); REGSTACK.pop_back();
        REGSTACK.push_back(op1 | op2);
        PC += 1;
    }},
    {"NOT", 23, 0, []() { 
        if (REGSTACK.empty()) return;
        int64_t op1 = REGSTACK.back(); REGSTACK.pop_back();
        REGSTACK.push_back(~op1);
        PC += 1;
    }},
    {"EQ", 31, 0, []() { 
        if (REGSTACK.size() < 2) return;
        int64_t op2 = REGSTACK.back(); REGSTACK.pop_back();
        int64_t op1 = REGSTACK.back(); REGSTACK.pop_back();
        REGSTACK.push_back(op1 == op2 ? 1 : 0);
        PC += 1;
    }},
    {"GT", 32, 0, []() { 
        if (REGSTACK.size() < 2) return;
        int64_t op2 = REGSTACK.back(); REGSTACK.pop_back();
        int64_t op1 = REGSTACK.back(); REGSTACK.pop_back();
        REGSTACK.push_back(op1 > op2 ? 1 : 0);
        PC += 1;
    }},
    {"LT", 33, 0, []() { 
        if (REGSTACK.size() < 2) return;
        int64_t op2 = REGSTACK.back(); REGSTACK.pop_back();
        int64_t op1 = REGSTACK.back(); REGSTACK.pop_back();
        REGSTACK.push_back(op1 < op2 ? 1 : 0);
        PC += 1;
    }},
    {"LOAD", 41, 1, []() { 
        int addr = MEM[CODE_SEG_BEGIN + PC + 1];
        if (addr >= 0 && addr < DATA_SEG_LEN) {
            REGSTACK.push_back(MEM[DATA_SEG_BEGIN + addr]);
        } else {
            REGSTACK.push_back(0);
        }
        PC += 2;
    }},
    {"STORE", 42, 1, []() { 
        int addr = MEM[CODE_SEG_BEGIN + PC + 1];
        if (!REGSTACK.empty() && addr >= 0 && addr < DATA_SEG_LEN) {
            MEM[DATA_SEG_BEGIN + addr] = REGSTACK.back();
            REGSTACK.pop_back();
        }
        PC += 2;
    }},
    {"LOAD_INDIRECT", 43, 0, []() { 
        if (REGSTACK.empty()) return;
        int addr = REGSTACK.back(); REGSTACK.pop_back();
        if (addr >= 0 && addr < DATA_SEG_LEN) {
            REGSTACK.push_back(MEM[DATA_SEG_BEGIN + addr]);
        } else {
            REGSTACK.push_back(0);
        }
        PC += 1;
    }},
    {"STORE_INDIRECT", 44, 0, []() { 
        if (REGSTACK.size() < 2) return;
        int addr = REGSTACK.back(); REGSTACK.pop_back();
        int64_t val = REGSTACK.back(); REGSTACK.pop_back();
        if (addr >= 0 && addr < DATA_SEG_LEN) {
            MEM[DATA_SEG_BEGIN + addr] = val;
        }
        PC += 1;
    }},
    {"JUMP", 51, 1, []() { 
        PC = MEM[CODE_SEG_BEGIN + PC + 1];
    }},
    {"JUMPIF", 52, 1, []() { 
        if (REGSTACK.empty()) {
            PC += 2;
        } else {
            int64_t condition = REGSTACK.back(); REGSTACK.pop_back();
            if (condition == 0) {
                PC += 2;
            } else {
                PC = MEM[CODE_SEG_BEGIN + PC + 1];
            }
        }
    }},
    {"PRINT", 61, 0, []() { 
        if (!REGSTACK.empty()) {
            cout << REGSTACK.back() << endl;
            REGSTACK.pop_back();
        }
        PC += 1;
    }},
    {"HALT", 255, 0, []() {}}
};

map<string, string> ALIAS_MAP = {
    {"+", "ADD"}, {"-", "SUB"}, {"*", "MUL"}, {"/", "DIV"}, {"%", "MOD"},
    {"=", "EQ"}, {">", "GT"}, {"<", "LT"},
    {"&", "AND"}, {"|", "OR"}, {"~", "NOT"},
    {"JMP", "JUMP"}, {"HLT", "HALT"}
};

// 辅助函数
bool isTag(const string& s) {
    return !s.empty() && s.back() == ':';
}

string tagName(const string& s) {
    if (s.empty()) return s;
    return s.substr(0, s.size() - 1);
}

Instruction* findInstruction(const string& name) {
    for (auto& instr : ISA) {
        if (instr.name == name) {
            return &instr;
        }
    }
    return nullptr;
}

// 汇编器 - 重新实现以匹配Racket版本
vector<int64_t> ASSEMBLE(const vector<string>& tokens) {
    map<string, int> tags;
    vector<int64_t> result;
    
    // 第一遍：收集标签
    int addr = 0;
    for (size_t i = 0; i < tokens.size(); i++) {
        const string& token = tokens[i];
        if (isTag(token)) {
            tags[tagName(token)] = addr;
        } else {
            // 检查是否为指令或别名
            string actual_name = token;
            if (ALIAS_MAP.find(token) != ALIAS_MAP.end()) {
                actual_name = ALIAS_MAP[token];
            }
            
            if (Instruction* instr = findInstruction(actual_name)) {
                addr++; // 指令码占用一个位置
                if (instr->args > 0) {
                    addr++; // 参数占用一个位置
                    i++; // 跳过参数
                }
            } else if (!token.empty() && (isdigit(token[0]) || token[0] == '-')) {
                addr++; // 数字占用一个位置
            }
            // 标签引用不增加地址计数
        }
    }
    
    // 第二遍：生成代码
    addr = 0;
    for (size_t i = 0; i < tokens.size(); i++) {
        const string& token = tokens[i];
        if (isTag(token)) {
            continue;
        }
        
        // 检查是否为指令或别名
        string actual_name = token;
        if (ALIAS_MAP.find(token) != ALIAS_MAP.end()) {
            actual_name = ALIAS_MAP[token];
        }
        
        if (Instruction* instr = findInstruction(actual_name)) {
            result.push_back(instr->code);
            addr++;
            
            if (instr->args > 0) {
                i++;
                if (i >= tokens.size()) {
                    result.push_back(0);
                } else {
                    const string& arg_token = tokens[i];
                    if (tags.find(arg_token) != tags.end()) {
                        result.push_back(tags[arg_token]);
                    } else {
                        try {
                            result.push_back(stoi(arg_token));
                        } catch (...) {
                            result.push_back(0);
                        }
                    }
                }
                addr++;
            }
        } else if (!token.empty() && (isdigit(token[0]) || token[0] == '-')) {
            try {
                result.push_back(stoi(token));
            } catch (...) {
                result.push_back(0);
            }
            addr++;
        } else if (tags.find(token) != tags.end()) {
            result.push_back(tags[token]);
            addr++;
        }
    }
    
    return result;
}

// 加载器
void LOADER(const vector<int64_t>& bytecode) {
    if (bytecode.size() > CODE_SEG_LEN) {
        throw runtime_error("Bytecode too large for code segment");
    }
    for (size_t i = 0; i < bytecode.size() && i < CODE_SEG_LEN; i++) {
        MEM[CODE_SEG_BEGIN + i] = bytecode[i];
    }
}

string format_memory_hexdump() {
    stringstream ss;
    ss << "#(";
    for (int i = 0; i < MEM_LEN; i++) {
        if (i > 0) ss << " ";
        ss << MEM[i];
    }
    ss << ")";
    return ss.str();
}

// 调试函数
void dump_machine() {
    cout << "((CODE-SEG-BEGIN " << CODE_SEG_BEGIN << ")" << endl;
    cout << " (CODE-SEG-LEN " << CODE_SEG_LEN << ")" << endl;
    cout << " (DATA-SEG-BEGIN " << DATA_SEG_BEGIN << ")" << endl;
    cout << " (DATA-SEG-LEN " << DATA_SEG_LEN << ")" << endl;
    cout << " (PC " << PC << ")" << endl;
    cout << " (REGSTACK (";
    for (size_t i = 0; i < REGSTACK.size(); i++) {
        if (i > 0) cout << " ";
        cout << REGSTACK[i];
    }
    cout << "))" << endl;
    cout << " (MEM " << format_memory_hexdump() << ")" << endl;
    cout << " (COINS-REMAIN " << COINS_REMAIN << ")" << endl;
    cout << " (CYCLES " << CYCLES << "))" << endl;
}

void debug_print() {
    map<int, string> dasm_map;
    for (const auto& instr : ISA) {
        dasm_map[instr.code] = instr.name;
    }
    
    int instr_code = MEM[CODE_SEG_BEGIN + PC];
    string instr_name = dasm_map.count(instr_code) ? dasm_map[instr_code] : "UNKNOWN";
    
    cout << "CYCLES: " << CYCLES << " REGSTACK: [";
    for (size_t i = 0; i < REGSTACK.size(); i++) {
        if (i > 0) cout << ", ";
        cout << REGSTACK[i];
    }
    cout << "] instr: " << instr_name;
    
    auto instr_it = find_if(ISA.begin(), ISA.end(), [instr_code](const Instruction& i) {
        return i.code == instr_code;
    });
    
    if (instr_it != ISA.end() && instr_it->args > 0) {
        cout << " " << MEM[CODE_SEG_BEGIN + PC + 1];
    }
    cout << endl;
}

bool step_machine() {
    debug_print();
    
    int instr_code = MEM[CODE_SEG_BEGIN + PC];
    auto instr_it = find_if(ISA.begin(), ISA.end(), [instr_code](const Instruction& i) {
        return i.code == instr_code;
    });
    
    if (instr_it != ISA.end()) {
        instr_it->action();
    } else {
        PC += 1;
    }
    
    COINS_REMAIN--;
    return instr_code != 255; // 如果不是HALT指令，返回true
}

void RUN_MACHINE() {
    while (COINS_REMAIN > 0 && PC >= 0 && PC < MEM_LEN) {
        if (!step_machine()) {
            break; // 遇到HALT指令
        }
        CYCLES++;
        
        if (PC >= MEM_LEN) {
            PC = 0;
        }
        
    }
    
    if (COINS_REMAIN <= 0) {
        cout << "NOT-ENOUGH-COINS" << endl;
    }
    dump_machine();
}

int main() {
    try {
        // 测试代码 - 与Racket版本完全一致
        vector<string> test_code = {
            "PUSH", "100",
            "STORE", "100",
            "LOOP:",
            "LOAD", "100",
            "PUSH", "1",
            "SUB",
            "STORE", "100",
            "LOAD", "100",
            "LOAD", "100",
            "STORE_INDIRECT",
            "LOAD", "100",
            "PUSH", "0",
            "EQ",
            "JUMPIF", "PRIME_LEFT",
            "JUMP", "LOOP",
            "PRIME_LEFT:",
            "PUSH", "2",
            "STORE", "1",
            "P_LOOP:",
            "LOAD", "1",
            "LOAD", "1",
            "MUL",
            "PUSH", "100",
            "GT",
            "JUMPIF", "HALT",
            "LOAD", "1",
            "LOAD_INDIRECT",
            "PUSH", "0",
            "EQ",
            "JUMPIF", "SKIP_MARK",
            "JUMP", "MARK_MULTIPLES",
            "MARK_MULTIPLES:",
            "LOAD", "1",
            "LOAD", "1",
            "MUL",
            "STORE", "0",
            "I_LOOP:",
            "LOAD", "0",
            "PUSH", "100",
            "LT",
            "JUMPIF", "I_CONTINUE",
            "JUMP", "SKIP_MARK",
            "I_CONTINUE:",
            "PUSH", "0",
            "LOAD", "0",
            "STORE_INDIRECT",
            "LOAD", "0",
            "LOAD", "1",
            "ADD",
            "STORE", "0",
            "JUMP", "I_LOOP",
            "SKIP_MARK:",
            "LOAD", "1",
            "PUSH", "1",
            "ADD",
            "STORE", "1",
            "JUMP", "P_LOOP",
            "HALT:",
            "HALT"
        };
        
        auto bytecode = ASSEMBLE(test_code);
        LOADER(bytecode);
        
        dump_machine();
        RUN_MACHINE();
        
    } catch (const exception& e) {
        cerr << "Error: " << e.what() << endl;
        return 1;
    }
    
    return 0;
}
